package main

import (
	"context"
	"fmt"
	"github.com/coreos/etcd/client"
	"github.com/coreos/etcd/pkg/transport"
	"log"
	"math"
	"time"
)

type ServiceReg struct {
	kApi        client.KeysAPI
	serviceList map[string]string
}

func NewServiceReg(addr []string, serverList map[string]string, timeout int64) (ser *ServiceReg, err error) {
	ser = &ServiceReg{serviceList: make(map[string]string, 16)}

	tr, err := transport.NewTransport(transport.TLSInfo{}, 5*time.Second)
	if err != nil {
		return
	}

	conf := client.Config{Endpoints: addr, Transport: tr}

	var c client.Client
	if c, err = client.New(conf); err != nil {
		return
	}

	ser.kApi = client.NewKeysAPI(c)

	for key, val := range serverList {
		_, err = ser.kApi.Set(context.Background(), key, val, nil)
		if err != nil {
			return
		}

		ser.serviceList[key] = val
	}

	go func() {
		o := client.SetOptions{
			TTL:       time.Second * time.Duration(timeout),
			PrevExist: client.PrevExist,
		}

		for {
			for s, addr := range ser.serviceList {
				_, err := ser.kApi.Set(context.Background(), s, addr, &o)
				if etErr, ok := err.(client.Error); ok && etErr.Code == client.ErrorCodeKeyNotFound {
					fmt.Printf("已经关闭续租功能\n")
					return
				} else if err != nil {
					fmt.Printf("续租失败: %v\n", err)
				} else {
					fmt.Printf("续租成功\n")
				}
			}

			time.Sleep(time.Second * time.Duration(math.Ceil(float64(timeout)*0.3)))
		}
	}()

	return
}

// 关闭租约服务
func (this *ServiceReg) Close() {
	for key, _ := range this.serviceList {
		_, _ = this.kApi.Delete(context.Background(), key, nil)
	}
}

// 使用方法
func main() {
	etcdAddr := []string{"127.0.0.1:2380"}
	servAddr := map[string]string{
		"/node/192.168.200.1": "192.168.200.1:50051",
		"/node/192.168.200.2": "192.168.200.2:50051",
	}
	ser, err := NewServiceReg(etcdAddr, servAddr, 10)
	if err != nil {
		log.Fatal(err)
	}
	defer ser.Close()
}
